import cv2
import os
import shutil
from pathlib import Path

class YOLOv5LabelingTool:
    def __init__(self):
        self.labels = {}
        self.current_image = None
        self.current_image_path = ""
        self.image_folder_path = ""
        self.image_save_path = ""
        self.label_save_path = ""
        self.current_boxes = []
        self.temp_point = None
        self.current_label_id = None
        self.window_name = "YOLOv5 标注工具"
        self.image_files = []
        self.current_image_index = 0
        
        # 为不同标签定义不同颜色 (BGR格式)
        self.label_colors = [
            (0, 255, 0),    # 绿色 - 标签0
            (255, 0, 0),    # 蓝色 - 标签1
            (0, 0, 255),    # 红色 - 标签2
            (255, 255, 0),  # 青色 - 标签3
            (255, 0, 255),  # 洋红色 - 标签4
            (0, 255, 255),  # 黄色 - 标签5
            (128, 0, 128),  # 紫色 - 标签6
            (255, 165, 0),  # 橙色 - 标签7
            (0, 128, 255),  # 浅蓝色 - 标签8
            (128, 255, 0),  # 浅绿色 - 标签9
        ]
        
        self.color_names = ["绿色", "蓝色", "红色", "青色", "洋红", "黄色", "紫色", "橙色", "浅蓝", "浅绿"]
    
    def setup_labels(self):
        """设置标签ID和名称"""
        print("=" * 60)
        print("YOLOv5 标注工具 - 标签设置")
        print("=" * 60)
        print("请输入标签信息 (输入 'done' 完成设置)")
        print("格式: ID号 标签名称")
        print("例如: 0 person")
        print("-" * 60)
        
        while True:
            user_input = input("请输入标签 (ID 名称): ").strip()
            if user_input.lower() == 'done':
                break
                
            try:
                parts = user_input.split(' ', 1)
                if len(parts) != 2:
                    print("❌ 格式错误，请使用: ID 名称")
                    continue
                    
                label_id = int(parts[0])
                label_name = parts[1]
                self.labels[label_id] = label_name
                color_idx = label_id % len(self.color_names)
                print(f"✅ 已添加标签: {label_id} -> {label_name} ({self.color_names[color_idx]})")
                
            except ValueError:
                print("❌ ID必须是数字")
                continue
        
        if not self.labels:
            print("❌ 未设置任何标签")
            return False
            
        print("\n当前标签列表:")
        for label_id, label_name in sorted(self.labels.items()):
            color_idx = label_id % len(self.color_names)
            print(f"  {label_id}: {label_name} ({self.color_names[color_idx]})")
        print("-" * 60)
        return True
    
    def setup_paths(self):
        """设置路径"""
        print("\n路径设置:")
        
        # 图片文件夹路径
        while True:
            self.image_folder_path = input("请输入图片文件夹路径: ").strip()
            if os.path.exists(self.image_folder_path) and os.path.isdir(self.image_folder_path):
                if self._load_image_files():
                    break
                else:
                    print("❌ 文件夹中没有找到有效的数字命名图片文件")
            else:
                print("❌ 文件夹不存在或不是有效目录")
        
        # 图片保存路径
        self.image_save_path = input("请输入图片保存路径: ").strip()
        os.makedirs(self.image_save_path, exist_ok=True)
        
        # 标签保存路径
        self.label_save_path = input("请输入标签保存路径: ").strip()
        os.makedirs(self.label_save_path, exist_ok=True)
        
        # 设置开始图片编号
        while True:
            try:
                start_num = input(f"请输入从第几张图片开始标注 (1-{len(self.image_files)}): ").strip()
                self.current_image_index = int(start_num) - 1
                if 0 <= self.current_image_index < len(self.image_files):
                    break
                else:
                    print(f"❌ 请输入 1 到 {len(self.image_files)} 之间的数字")
            except ValueError:
                print("❌ 请输入有效的数字")
        
        print(f"✅ 设置完成，将从第{self.current_image_index + 1}张图片开始")
        return True
    
    def _load_image_files(self):
        """加载图片文件列表"""
        valid_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']
        image_files = []
        
        for file in os.listdir(self.image_folder_path):
            file_ext = os.path.splitext(file)[1].lower()
            if file_ext in valid_extensions:
                file_name = os.path.splitext(file)[0]
                try:
                    int(file_name)
                    image_files.append(file)
                except ValueError:
                    continue
        
        if not image_files:
            return False
        
        # 按数字排序
        image_files.sort(key=lambda x: int(os.path.splitext(x)[0]))
        self.image_files = image_files
        
        print(f"✅ 找到 {len(image_files)} 个图片文件")
        return True
    
    def load_current_image(self):
        """加载当前图片"""
        if self.current_image_index >= len(self.image_files):
            return False
            
        current_file = self.image_files[self.current_image_index]
        image_path = os.path.join(self.image_folder_path, current_file)
        
        self.current_image = cv2.imread(image_path)
        if self.current_image is not None:
            self.current_image_path = image_path
            print(f"✅ 已加载: {current_file} (第{self.current_image_index + 1}/{len(self.image_files)}张)")
            return True
        return False
    
    def mouse_callback(self, event, x, y, flags, param):
        """鼠标回调函数"""
        if event == cv2.EVENT_LBUTTONDOWN and self.current_label_id is not None:
            if self.current_image is None:
                return
                
            if self.temp_point is None:
                self.temp_point = (x, y)
                print(f"左上角点: ({x}, {y})")
                self.draw_image()
            else:
                x1, y1 = self.temp_point
                left, top = min(x1, x), min(y1, y)
                right, bottom = max(x1, x), max(y1, y)
                
                # 检查矩形大小
                if abs(right - left) < 10 or abs(bottom - top) < 10:
                    print("⚠️  矩形太小，请重新选择")
                    self.temp_point = None
                    self.draw_image()
                    return
                
                # 转换为YOLO格式
                img_h, img_w = self.current_image.shape[:2]
                center_x = (left + right) / 2.0 / img_w
                center_y = (top + bottom) / 2.0 / img_h
                width = (right - left) / img_w
                height = (bottom - top) / img_h
                
                box_info = {
                    'label_id': self.current_label_id,
                    'center_x': center_x,
                    'center_y': center_y,
                    'width': width,
                    'height': height,
                    'pixel_coords': (left, top, right, bottom)
                }
                self.current_boxes.append(box_info)
                
                label_name = self.labels[self.current_label_id]
                color_name = self.color_names[self.current_label_id % len(self.color_names)]
                print(f"✅ 添加标注框: 标签{self.current_label_id}({label_name},{color_name}) 坐标({left},{top},{right},{bottom})")
                self.temp_point = None
                self.draw_image()
    
    def draw_image(self):
        """绘制图片和标注框"""
        if self.current_image is None:
            return
            
        display_img = self.current_image.copy()
        
        # 绘制标注框
        for box in self.current_boxes:
            left, top, right, bottom = box['pixel_coords']
            label_id = box['label_id']
            label_name = self.labels[label_id]
            
            # 获取标签对应的颜色
            color = self.label_colors[label_id % len(self.label_colors)]
            
            # 绘制矩形框
            cv2.rectangle(display_img, (left, top), (right, bottom), color, 3)
            
            # 绘制标签背景
            label_text = f"{label_id}: {label_name}"
            (text_width, text_height), _ = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
            cv2.rectangle(display_img, (left, top-text_height-15), (left+text_width+10, top), color, -1)
            
            # 绘制标签文本
            cv2.putText(display_img, label_text, (left+5, top-8), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
        
        # 绘制临时点
        if self.temp_point is not None:
            cv2.circle(display_img, self.temp_point, 8, (0, 0, 255), -1)
            cv2.circle(display_img, self.temp_point, 12, (0, 0, 255), 2)
            
            # 显示当前选择的标签信息
            if self.current_label_id is not None:
                color_name = self.color_names[self.current_label_id % len(self.color_names)]
                temp_text = f"选择标签: {self.current_label_id} ({self.labels[self.current_label_id]},{color_name})"
                cv2.putText(display_img, temp_text, (10, 35), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
        
        # 显示图片信息
        info_text = f"图片: {self.current_image_index + 1}/{len(self.image_files)} | 标注框: {len(self.current_boxes)}"
        cv2.putText(display_img, info_text, (10, display_img.shape[0] - 15), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
        
        # 显示操作提示
        if self.current_label_id is not None and self.temp_point is None:
            hint_text = "点击图片选择左上角"
            cv2.putText(display_img, hint_text, (10, 70), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
        elif self.current_label_id is not None and self.temp_point is not None:
            hint_text = "点击图片选择右下角"
            cv2.putText(display_img, hint_text, (10, 70), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
            
        cv2.imshow(self.window_name, display_img)
    
    def save_labels(self):
        """保存标签文件"""
        image_name = Path(self.current_image_path).stem
        label_file_path = os.path.join(self.label_save_path, f"{image_name}.txt")
        
        try:
            with open(label_file_path, 'w') as f:
                for box in self.current_boxes:
                    line = f"{box['label_id']} {box['center_x']:.6f} {box['center_y']:.6f} {box['width']:.6f} {box['height']:.6f}\n"
                    f.write(line)
            print(f"✅ 已保存标签文件: {image_name}.txt (包含{len(self.current_boxes)}个标注框)")
        except Exception as e:
            print(f"❌ 保存失败: {e}")
    
    def copy_image(self):
        """复制图片到保存目录"""
        try:
            image_name = os.path.basename(self.current_image_path)
            dest_path = os.path.join(self.image_save_path, image_name)
            if self.current_image_path != dest_path:
                shutil.copy2(self.current_image_path, dest_path)
                print(f"✅ 已复制图片: {image_name}")
        except Exception as e:
            print(f"❌ 复制图片失败: {e}")

    def start_labeling(self):
        """开始标注"""
        cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
        cv2.setMouseCallback(self.window_name, self.mouse_callback)

        print("\n" + "=" * 60)
        print("标注说明:")
        print("1. 输入标签ID选择标注类别")
        print("2. 在图片上点击两次确定矩形框")
        print("3. 输入 'next' 保存并切换到下一张图片")
        print("4. 输入 'save' 保存当前标注")
        print("5. 输入 'clear' 清除当前所有标注")
        print("6. 输入 'quit' 退出程序")
        print("7. 按ESC键也可以退出")
        print("=" * 60)

        while self.current_image_index < len(self.image_files):
            if not self.load_current_image():
                break

            self.current_boxes = []
            self.temp_point = None
            self.current_label_id = None
            self.draw_image()

            while True:
                print(f"\n当前图片: {os.path.basename(self.current_image_path)} "
                      f"(第{self.current_image_index + 1}/{len(self.image_files)}张)")
                print(f"标注框数量: {len(self.current_boxes)}")
                print("标签列表:")
                for label_id, label_name in sorted(self.labels.items()):
                    color_idx = label_id % len(self.color_names)
                    print(f"  {label_id}: {label_name} ({self.color_names[color_idx]})")

                # 处理OpenCV窗口事件
                key = cv2.waitKey(1) & 0xFF
                if key == 27:  # ESC键退出
                    cv2.destroyAllWindows()
                    return

                command = input("\n请输入命令: ").strip()

                if command.lower() == 'quit':
                    cv2.destroyAllWindows()
                    return
                elif command.lower() == 'next':
                    self.save_labels()
                    self.copy_image()
                    self.current_image_index += 1
                    break
                elif command.lower() == 'save':
                    self.save_labels()
                    self.copy_image()
                elif command.lower() == 'clear':
                    self.current_boxes = []
                    self.temp_point = None
                    self.current_label_id = None
                    self.draw_image()
                    print("✅ 已清除所有标注")
                else:
                    try:
                        label_id = int(command)
                        if label_id in self.labels:
                            self.current_label_id = label_id
                            self.temp_point = None
                            self.draw_image()
                            color_name = self.color_names[label_id % len(self.color_names)]
                            print(f"✅ 当前标注标签: {label_id} ({self.labels[label_id]},{color_name})")
                            print("请在图片上点击两次确定矩形框")
                        else:
                            print(f"❌ 标签ID {label_id} 不存在")
                    except ValueError:
                        print("❌ 无效命令")

        print("✅ 所有图片标注完成！")
        cv2.destroyAllWindows()

    def run(self):
        """运行标注工具"""
        print("=" * 60)
        print("YOLOv5 标注工具")
        print("=" * 60)

        if not self.setup_labels():
            return
        if not self.setup_paths():
            return

        self.start_labeling()
        print("\n✅ 标注工具已退出")

def main():
    tool = YOLOv5LabelingTool()
    tool.run()

if __name__ == "__main__":
    main()
